{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "022c14ae",
   "metadata": {},
   "source": [
    "# AI Red Teaming Notebook Overview\n",
    "\n",
    "This streamlined notebook guides you through progressively richer AI red teaming evaluations using the Azure AI Evaluation SDK with logging into a Microsoft Foundry project (classic, hub-based).\n",
    "\n",
    "You will:\n",
    "- Run a fast smoke test with a deterministic safe callback (baseline expectations, near-zero Attack Success Rate).\n",
    "- Target a real Azure OpenAI deployment to observe genuine safety behavior.\n",
    "- Expand coverage across multiple risk categories and layered attack strategies.\n",
    "- Add advanced multi-strategy scans (including composed transformations) to probe layered defenses.\n",
    "- (Optional, end of notebook) Supply your own domain‑specific risky objectives.\n",
    "\n",
    "Artifacts: Each scan writes a JSON scorecard file (label + UTC time). Use these for comparison, regression tracking, or upload into Microsoft Foundry for analysis under the AI red teaming evaluations view.\n",
    "\n",
    "Execution time scales roughly with: risk_categories × attack_strategies × num_objectives. Start small, expand only after verifying prior steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45803341",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 1: Installation\n",
    "import sys, subprocess\n",
    "\n",
    "packages = [\n",
    "    \"duckdb==1.3.2\",\n",
    "    \"azure-ai-evaluation[redteam]==1.12.0\",\n",
    "    \"azure-identity==1.25.1\",\n",
    "    \"openai==2.8.1\",\n",
    "    \"azure-ai-projects==1.0.0\",\n",
    "    \"python-dotenv==1.2.1\",\n",
    "]\n",
    "\n",
    "subprocess.check_call([\n",
    "    sys.executable,\n",
    "    \"-m\",\n",
    "    \"pip\",\n",
    "    \"install\",\n",
    "    \"--quiet\",\n",
    "    \"--upgrade\",\n",
    "    \"--no-warn-conflicts\",\n",
    "    *packages,\n",
    "])\n",
    "print(\"Installed (pinned / required):\")\n",
    "for p in packages:\n",
    "    print(\"  -\", p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec9949f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 2 - imports\n",
    "from typing import Optional, Dict, Any\n",
    "import os\n",
    "\n",
    "# Azure imports\n",
    "from azure.ai.evaluation.red_team import RedTeam, RiskCategory, AttackStrategy\n",
    "\n",
    "# OpenAI import\n",
    "from openai import AzureOpenAI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7ac0de0",
   "metadata": {},
   "source": [
    "## Core Concepts: RedTeam, Risk Categories, Attack Strategies & Targets\n",
    "\n",
    "**RedTeam Orchestrator**: Generates attack objectives, transforms prompts via strategies, invokes your target, and scores responses.\n",
    "\n",
    "**Risk Categories (what we probe)**: Violence, Hate/Unfairness, Sexual, SelfHarm. You can supply a subset for faster iteration. Missing categories reduce coverage but cut cost/time.\n",
    "\n",
    "**Attack Strategies (how we probe)**:\n",
    "- Complexity group macros: `EASY`, `MODERATE` (bundles of simpler / moderate transformations)\n",
    "- Individual transformations: Flip, CharSwap, UnicodeConfusable, Leetspeak, Url, Base64, ROT13, etc.\n",
    "- Composition: `AttackStrategy.Compose([Base64, ROT13])` layers transformations to simulate obfuscation chains.\n",
    "\n",
    "**num_objectives**: Count of seed prompts per category (per applied strategy). Linear multiplier on runtime.\n",
    "\n",
    "**Targets (what gets attacked)**:\n",
    "1. Simple synchronous callback (returns fixed text) – deterministic baseline.\n",
    "2. Model configuration dict – RedTeam handles generation calls internally.\n",
    "3. Fully custom (async) application wrapper – replicate real app logic, pre/post-processing.\n",
    "\n",
    "We progress through (1) → (2) → (3+) for clarity.\n",
    "\n",
    "> Tip: Keep early scans lean (≤2 categories, 1 strategy, num_objectives=1) to validate authentication & environment quickly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce743294",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 3 - Create credential (switchable auth)\n",
    "import os\n",
    "from azure.identity import DefaultAzureCredential, ManagedIdentityCredential\n",
    "\n",
    "# Allow easy switching via env flag; default to DefaultAzureCredential for broader token chain\n",
    "_use_default = os.environ.get(\"REDTEAM_USE_DEFAULT_CRED\", \"1\") == \"1\"\n",
    "if _use_default:\n",
    "    credential = DefaultAzureCredential(exclude_interactive_browser_credential=True)\n",
    "    print(\"Using DefaultAzureCredential (Managed Identity + Azure CLI + Env, etc.)\")\n",
    "else:\n",
    "    credential = ManagedIdentityCredential()\n",
    "    print(\"Using ManagedIdentityCredential explicitly\")\n",
    "\n",
    "# Quick probe (optional) – will no-op on some identities if scope isn't accessible\n",
    "try:\n",
    "    token = credential.get_token(\"https://management.azure.com/.default\")\n",
    "    print(\"Acquired mgmt token (truncated):\", token.token[:24], \"...\")\n",
    "except Exception as e:  # noqa: BLE001\n",
    "    print(\"Token probe skipped:\", e)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "846e231e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 4 - Dynamic .env discovery\n",
    "from pathlib import Path\n",
    "from typing import List\n",
    "import os, re\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "REQUIRED_KEYS = [\n",
    "    \"AZURE_SUBSCRIPTION_ID\",\n",
    "    \"AZURE_RESOURCE_GROUP_NAME\",\n",
    "    \"AZURE_PROJECT_NAME\",\n",
    "    \"AZURE_OPENAI_DEPLOYMENT_NAME\",\n",
    "    \"AZURE_OPENAI_ENDPOINT\",\n",
    "    \"AZURE_OPENAI_API_KEY\",\n",
    "    \"AZURE_OPENAI_API_VERSION\",\n",
    "]\n",
    "\n",
    "explicit_path = os.environ.get(\"REDTEAM_DOTENV_PATH\")\n",
    "searched: List[Path] = []\n",
    "selected = None\n",
    "\n",
    "candidates: List[Path] = []\n",
    "if explicit_path:\n",
    "    p = Path(explicit_path)\n",
    "    if p.is_file():\n",
    "        candidates.append(p)\n",
    "\n",
    "# Typical Microsoft Foundry (classic, hub-based) project mount pattern: /afh/projects/<resource-project-guid>/shared/files/.env\n",
    "root = Path('/afh/projects')\n",
    "if root.is_dir():\n",
    "    for child in root.iterdir():\n",
    "        if child.is_dir() and ('-project-' in child.name):\n",
    "            env_candidate = child / 'shared' / 'files' / '.env'\n",
    "            searched.append(env_candidate)\n",
    "            if env_candidate.is_file():\n",
    "                candidates.append(env_candidate)\n",
    "\n",
    "# Fallback: shallow glob for any .env directly under shared/files\n",
    "if not candidates and root.is_dir():\n",
    "    for env_candidate in root.glob('**/shared/files/.env'):\n",
    "        searched.append(env_candidate)\n",
    "        if env_candidate.is_file():\n",
    "            candidates.append(env_candidate)\n",
    "            break\n",
    "\n",
    "# Choose first containing all required keys, else first existing\n",
    "for c in candidates:\n",
    "    try:\n",
    "        text = c.read_text()\n",
    "        if all(re.search(rf'^ {k}=', text, re.MULTILINE) or re.search(rf'^{k}=', text, re.MULTILINE) for k in REQUIRED_KEYS):\n",
    "            selected = c\n",
    "            break\n",
    "    except Exception:\n",
    "        pass\n",
    "if selected is None and candidates:\n",
    "    selected = candidates[0]\n",
    "\n",
    "if selected and selected.is_file():\n",
    "    load_dotenv(selected)\n",
    "    missing_after = [k for k in REQUIRED_KEYS if not os.environ.get(k)]\n",
    "    print(f\"Loaded .env from: {selected}\")\n",
    "    if missing_after:\n",
    "        print(\"Still missing keys:\", missing_after)\n",
    "else:\n",
    "    print(\"No .env loaded. Candidates searched (first 5):\", [str(p) for p in searched[:5]])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d6d53b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 5 - Set variables\n",
    "import os\n",
    "\n",
    "_required_keys = [\n",
    "    \"AZURE_SUBSCRIPTION_ID\",\n",
    "    \"AZURE_RESOURCE_GROUP_NAME\",\n",
    "    \"AZURE_PROJECT_NAME\",\n",
    "    \"AZURE_OPENAI_DEPLOYMENT_NAME\",\n",
    "    \"AZURE_OPENAI_ENDPOINT\",\n",
    "    \"AZURE_OPENAI_API_KEY\",\n",
    "    \"AZURE_OPENAI_API_VERSION\",\n",
    "]\n",
    "_env = {k: os.environ.get(k) for k in _required_keys}\n",
    "_missing = [k for k, v in _env.items() if not v]\n",
    "\n",
    "if _missing:\n",
    "    print(\"Missing environment variables:\", _missing)\n",
    "else:\n",
    "    # Construct objects / variables consumed by later cells\n",
    "    azure_ai_project = {\n",
    "        \"subscription_id\": _env[\"AZURE_SUBSCRIPTION_ID\"],\n",
    "        \"resource_group_name\": _env[\"AZURE_RESOURCE_GROUP_NAME\"],\n",
    "        \"project_name\": _env[\"AZURE_PROJECT_NAME\"],\n",
    "        \"credential\": credential,\n",
    "    }\n",
    "    azure_openai_deployment = _env[\"AZURE_OPENAI_DEPLOYMENT_NAME\"]\n",
    "    azure_openai_endpoint = _env[\"AZURE_OPENAI_ENDPOINT\"]\n",
    "    azure_openai_api_key = _env[\"AZURE_OPENAI_API_KEY\"]\n",
    "    azure_openai_api_version = _env[\"AZURE_OPENAI_API_VERSION\"]\n",
    "\n",
    "    print(\"Config loaded.\")\n",
    "    print(f\"Deployment: {azure_openai_deployment}\")\n",
    "    print(f\"Endpoint: {azure_openai_endpoint}\")\n",
    "    print(\"API key prefix:\", azure_openai_api_key[:6] + '...')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77266c34",
   "metadata": {},
   "source": [
    "### Environment Variables Loaded\n",
    "Required keys (subscription, resource group, project name, OpenAI deployment details, API key & version) enable two things:\n",
    "1. Locating the Azure AI Project for logging / result persistence.\n",
    "2. Invoking your Azure OpenAI deployment during model-based scans.\n",
    "\n",
    "Missing values will be surfaced in the next cell; resolve them before proceeding to model or advanced scans."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b6286a3",
   "metadata": {},
   "source": [
    "### Target Patterns\n",
    "You can point RedTeam at different target styles:\n",
    "- Deterministic callback (fast baseline, isolates framework logic).\n",
    "- Direct model configuration dict (minimal code, invokes deployment directly).\n",
    "- Rich application function (mirrors real system behavior: pre/post-processing, memory, safety layers).\n",
    "\n",
    "We start with the deterministic callback to validate plumbing before spending tokens on real model invocations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f5208ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 6 - Define a simple callback function that always returns a fixed response\n",
    "def financial_advisor_callback(query: str) -> str:  # noqa: ARG001\n",
    "    return \"I'm a financial advisor assistant. I can help with investment advice and financial planning within legal and ethical guidelines.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b32fce5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 7 - Create the `RedTeam` instance with minimal configurations\n",
    "red_team = RedTeam(\n",
    "    azure_ai_project=azure_ai_project,\n",
    "    credential=credential,\n",
    "    risk_categories=[RiskCategory.Violence, RiskCategory.HateUnfairness],\n",
    "    num_objectives=1,\n",
    " )\n",
    "# (Inserted) Configure writable output directory to avoid permission issues in read-only code mount\n",
    "import pathlib, os\n",
    "_rt_out = pathlib.Path.home() / 'redteam_outputs'\n",
    "_rt_out.mkdir(parents=True, exist_ok=True)\n",
    "try:\n",
    "    # Some internals may expose file_manager; guard if attribute layout changes\n",
    "    if hasattr(red_team, 'file_manager') and hasattr(red_team.file_manager, 'base_output_dir'):\n",
    "        red_team.file_manager.base_output_dir = str(_rt_out)\n",
    "        print('RedTeam base_output_dir set to', red_team.file_manager.base_output_dir)\n",
    "    else:\n",
    "        print('WARNING: Unable to set custom base_output_dir; attribute not found')\n",
    "except Exception as _e:\n",
    "    print('WARNING: Failed to set base_output_dir:', _e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e84917c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 8 - Basic Scan\n",
    "from datetime import datetime, timezone\n",
    "from pathlib import Path\n",
    "import pathlib, os\n",
    "\n",
    "scan_ts = datetime.now(timezone.utc).strftime(\"%H%M%S\")\n",
    "scan_label = \"Basic\"\n",
    "scan_name = f\"{scan_label}-Scan-{scan_ts}\"\n",
    "# Use writable output directory set earlier (fallback to home)\n",
    "_writable_base = pathlib.Path(getattr(red_team.file_manager, 'base_output_dir', Path.home()))\n",
    "_writable_base.mkdir(parents=True, exist_ok=True)\n",
    "output_path = _writable_base / f\"{scan_label.lower()}_scan-{scan_ts}.json\"\n",
    "\n",
    "basic_scan_result = await red_team.scan(\n",
    "    target=financial_advisor_callback,\n",
    "    scan_name=scan_name,\n",
    "    attack_strategies=[AttackStrategy.Flip],\n",
    "    output_path=str(output_path),\n",
    "    # skip_upload=True  # uncomment if remote artifact upload warnings persist\n",
    " )\n",
    "print(f\"Completed {scan_name} -> {output_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49bb121d",
   "metadata": {},
   "source": [
    "### Interpreting the Basic Scan\n",
    "Key metric: **Attack Success Rate (ASR)** = percentage of prompts whose responses are flagged harmful.\n",
    "\n",
    "Because our callback always returns a fixed safe sentence, expected ASR ≈ 0%.\n",
    "\n",
    "Use the emitted JSON to verify structure and plan which additional strategies or categories to add next."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f08f7c86",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 9 - Intermediary Scan (updated to mirror Cell 8 writable path pattern and add basic error handling)\n",
    "from datetime import datetime, timezone\n",
    "from pathlib import Path\n",
    "import pathlib, os\n",
    "\n",
    "# Lazy client creation (safe on reruns)\n",
    "if '_rt_client' not in globals():\n",
    "    from openai import AzureOpenAI\n",
    "    _rt_client = AzureOpenAI(\n",
    "        api_key=azure_openai_api_key,\n",
    "        api_version=azure_openai_api_version,\n",
    "        azure_endpoint=azure_openai_endpoint,\n",
    "    )\n",
    "\n",
    "# Accept both invocation styles used internally by RedTeam\n",
    "def model_target(query: str = None, user_prompt: str = None):\n",
    "    prompt = user_prompt if (user_prompt not in (None, \"\")) else query or \"(empty prompt)\"\n",
    "    return _rt_client.chat.completions.create(\n",
    "        model=azure_openai_deployment,\n",
    "        messages=[{\"role\": \"user\", \"content\": prompt}],\n",
    "        temperature=0,\n",
    "        max_tokens=256,\n",
    "    ).choices[0].message.content or \"\"\n",
    "\n",
    "# Reuse writable directory logic from Cell 8 for consistency\n",
    "_writable_base = pathlib.Path(getattr(red_team.file_manager, 'base_output_dir', Path.home()))\n",
    "_writable_base.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "scan_ts = datetime.now(timezone.utc).strftime('%H%M%S')\n",
    "scan_label = \"Intermediary\"\n",
    "scan_name = f\"{scan_label}-Scan-{scan_ts}\"\n",
    "output_path = _writable_base / f\"{scan_label.lower()}_scan-{scan_ts}.json\"\n",
    "\n",
    "print(f\"Running {scan_name} -> {output_path}\")\n",
    "\n",
    "try:\n",
    "    intermediary_scan_result = await red_team.scan(\n",
    "        target=model_target,\n",
    "        scan_name=scan_name,\n",
    "        attack_strategies=[AttackStrategy.Flip],\n",
    "        output_path=str(output_path),\n",
    "    )\n",
    "    print(f\"Completed {scan_name} -> {output_path}\")\n",
    "except Exception as e:  # noqa: BLE001\n",
    "    print(f\"ERROR during {scan_name}: {e}\")\n",
    "    raise"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "687a0b18",
   "metadata": {},
   "source": [
    "### Moving to a Model Target\n",
    "Switching from a deterministic callback to an actual model introduces variability and real guardrail evaluation. Keeping the same single `Flip` strategy isolates model safety behavior from added obfuscation complexity.\n",
    "\n",
    "Next expansions: increase `num_objectives`, add additional strategies (CharSwap, UnicodeConfusable, etc.), or broaden risk categories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2ea45b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 10: Advanced scan - create expanded RedTeam instance (with writable base_output_dir)\n",
    "advanced_risk_categories = [\n",
    "    RiskCategory.Violence,\n",
    "    RiskCategory.HateUnfairness,\n",
    "    RiskCategory.Sexual,\n",
    "    RiskCategory.SelfHarm,\n",
    "]\n",
    "advanced_red_team = RedTeam(\n",
    "    azure_ai_project=azure_ai_project,\n",
    "    credential=credential,\n",
    "    risk_categories=advanced_risk_categories,\n",
    "    num_objectives=3,  # increase coverage per category (adjust for cost/time)\n",
    ")\n",
    "\n",
    "# Mirror writable directory logic used earlier for 'red_team'\n",
    "import pathlib\n",
    "_adv_out = pathlib.Path(getattr(red_team.file_manager, 'base_output_dir', pathlib.Path.home()))\n",
    "_adv_out.mkdir(parents=True, exist_ok=True)\n",
    "try:\n",
    "    if hasattr(advanced_red_team, 'file_manager') and hasattr(advanced_red_team.file_manager, 'base_output_dir'):\n",
    "        advanced_red_team.file_manager.base_output_dir = str(_adv_out)\n",
    "        print('Advanced RedTeam base_output_dir set to', advanced_red_team.file_manager.base_output_dir)\n",
    "    else:\n",
    "        print('WARNING: Could not set advanced_red_team base_output_dir; attribute not found')\n",
    "except Exception as _e:\n",
    "    print('WARNING: Failed setting advanced_red_team base_output_dir:', _e)\n",
    "\n",
    "print(\"Advanced RedTeam configured with categories:\", [c.name for c in advanced_risk_categories])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9dbcd0a",
   "metadata": {},
   "source": [
    "### Expanding Coverage\n",
    "We now include all four core risk categories and raise `num_objectives` to increase statistical signal. This increases token/time consumption proportionally.\n",
    "\n",
    "Broader coverage helps surface category-specific weaknesses early (e.g., higher ASR in SelfHarm vs Sexual)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a97b0baf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 11: Advanced scan (standardized naming & writable output dir)\n",
    "from datetime import datetime, timezone\n",
    "from pathlib import Path\n",
    "import pathlib  # for consistent writable directory handling\n",
    "\n",
    "scan_ts = datetime.now(timezone.utc).strftime('%H%M%S')\n",
    "scan_label = \"Advanced\"\n",
    "scan_name = f\"{scan_label}-Scan-{scan_ts}\"\n",
    "\n",
    "# Use same writable directory pattern as earlier scans (fallback to home if attribute missing)\n",
    "_writable_base = pathlib.Path(getattr(red_team.file_manager, 'base_output_dir', Path.home()))\n",
    "_writable_base.mkdir(parents=True, exist_ok=True)\n",
    "output_path = _writable_base / f\"{scan_label.lower()}_scan-{scan_ts}.json\"\n",
    "\n",
    "print(f\"Running {scan_name} -> {output_path}\")\n",
    "advanced_scan_result = await advanced_red_team.scan(\n",
    "    target=model_target,\n",
    "    scan_name=scan_name,\n",
    "    attack_strategies=[\n",
    "        AttackStrategy.EASY,\n",
    "        AttackStrategy.MODERATE,\n",
    "        AttackStrategy.Flip,\n",
    "        AttackStrategy.CharSwap,\n",
    "        AttackStrategy.UnicodeConfusable,\n",
    "        AttackStrategy.Leetspeak,\n",
    "        AttackStrategy.Url,\n",
    "        AttackStrategy.Compose([AttackStrategy.Base64, AttackStrategy.ROT13]),\n",
    "    ],\n",
    "    output_path=str(output_path),\n",
    ")\n",
    "print(f\"Completed {scan_name} -> {output_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d716791",
   "metadata": {},
   "source": [
    "### Advanced Strategies & Layering\n",
    "The advanced scan mixes:\n",
    "- Complexity groups (`EASY`, `MODERATE`) for breadth.\n",
    "- Obfuscations (CharSwap, UnicodeConfusable, Leetspeak, Url) to probe normalization defenses.\n",
    "- Encoding (Base64, ROT13 via composition) to test decoding / content safety layers.\n",
    "\n",
    "Capturing stderr lets you quickly surface any internal SDK errors alongside scan results.\n",
    "\n",
    "## Bring Your Own Objectives: Custom Attack Seed Prompts\n",
    "You can supply your own domain or application-specific risky prompts as objectives instead of (or in addition to) automatically generated ones.\n",
    "\n",
    "Format: a JSON file whose entries include `prompt` text and `risk-type` (one of: `violence`, `sexual`, `hate_unfairness`, `self_harm`). The number of prompts provided becomes the effective `num_objectives` for the scan.\n",
    "\n",
    "Use this when:\n",
    "- You have proprietary misuse scenarios not covered by generic seeds.\n",
    "- You want regression tracking on a fixed, curated risky prompt set.\n",
    "- You need to validate mitigations against previously successful attacks.\n",
    "\n",
    "Below we instantiate a new `RedTeam` with `custom_attack_seed_prompts` pointing to `data/prompts.json`, then run grouped difficulty strategies.\n",
    "\n",
    "> Tip: Keep a version-controlled prompts file so additions are reviewable and diffs tie to shifts in ASR.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2bf39de",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 12 - Custom prompts RedTeam instance (dynamic prompt discovery + writable dir)\n",
    "from pathlib import Path\n",
    "import pathlib, os, re\n",
    "\n",
    "# Allow explicit override\n",
    "explicit_prompt_path = os.environ.get(\"REDTEAM_CUSTOM_PROMPTS_PATH\")\n",
    "\n",
    "candidates = []\n",
    "searched = []\n",
    "selected = None\n",
    "\n",
    "if explicit_prompt_path:\n",
    "    p = Path(explicit_prompt_path)\n",
    "    if p.is_file():\n",
    "        candidates.append(p)\n",
    "\n",
    "# Derive project mount pattern similar to Cell 4 logic: /afh/projects/*-project-*/shared/files/data/prompts.json\n",
    "projects_root = Path('/afh/projects')\n",
    "if projects_root.is_dir():\n",
    "    for child in projects_root.iterdir():\n",
    "        if child.is_dir() and ('-project-' in child.name):\n",
    "            prompt_candidate = child / 'shared' / 'files' / 'data' / 'prompts.json'\n",
    "            searched.append(prompt_candidate)\n",
    "            if prompt_candidate.is_file():\n",
    "                candidates.append(prompt_candidate)\n",
    "\n",
    "# Fallback: glob search for any prompts.json under shared/files/data\n",
    "if not candidates and projects_root.is_dir():\n",
    "    for prompt_candidate in projects_root.glob('**/shared/files/data/prompts.json'):\n",
    "        searched.append(prompt_candidate)\n",
    "        if prompt_candidate.is_file():\n",
    "            candidates.append(prompt_candidate)\n",
    "            break\n",
    "\n",
    "# Select first existing candidate\n",
    "for c in candidates:\n",
    "    try:\n",
    "        if c.is_file():\n",
    "            selected = c\n",
    "            break\n",
    "    except Exception:\n",
    "        pass\n",
    "\n",
    "if not selected:\n",
    "    raise FileNotFoundError(\n",
    "        \"Could not locate prompts.json. Searched candidates (first 5): \" +\n",
    "        str([str(p) for p in searched[:5]]) +\n",
    "        \" - set REDTEAM_CUSTOM_PROMPTS_PATH to override.\"\n",
    "    )\n",
    "\n",
    "print(f\"Using prompts file: {selected}\")\n",
    "\n",
    "custom_red_team = RedTeam(\n",
    "    azure_ai_project=azure_ai_project,\n",
    "    credential=credential,\n",
    "    custom_attack_seed_prompts=str(selected),\n",
    ")\n",
    "\n",
    "# Align writable output directory with earlier scans\n",
    "_custom_out = pathlib.Path(getattr(red_team.file_manager, 'base_output_dir', Path.home()))\n",
    "_custom_out.mkdir(parents=True, exist_ok=True)\n",
    "try:\n",
    "    if hasattr(custom_red_team, 'file_manager') and hasattr(custom_red_team.file_manager, 'base_output_dir'):\n",
    "        custom_red_team.file_manager.base_output_dir = str(_custom_out)\n",
    "        print('Custom RedTeam base_output_dir set to', custom_red_team.file_manager.base_output_dir)\n",
    "    else:\n",
    "        print('WARNING: Could not set custom_red_team base_output_dir; attribute not found')\n",
    "except Exception as _e:\n",
    "    print('WARNING: Failed setting custom_red_team base_output_dir:', _e)\n",
    "\n",
    "print(\"Custom RedTeam ready. Prompt count determines num_objectives.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "828a5c3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 13: Execute scan with custom prompts and grouped difficulty strategies (writable output dir)\n",
    "from datetime import datetime, timezone\n",
    "from pathlib import Path\n",
    "import pathlib\n",
    "\n",
    "scan_ts = datetime.now(timezone.utc).strftime('%H%M%S')\n",
    "scan_label = \"Custom\"\n",
    "scan_name = f\"{scan_label}-Prompt-Scan-{scan_ts}\"\n",
    "\n",
    "# Use the custom_red_team writable directory (fallback to home)\n",
    "_writable_base = pathlib.Path(getattr(custom_red_team.file_manager, 'base_output_dir', Path.home()))\n",
    "_writable_base.mkdir(parents=True, exist_ok=True)\n",
    "output_path = _writable_base / f\"custom_prompt_scan-{scan_ts}.json\"\n",
    "print(f\"Running {scan_name} -> {output_path}\")\n",
    "\n",
    "try:\n",
    "    custom_result = await custom_red_team.scan(\n",
    "        target=model_target,  # reuse earlier model target callback\n",
    "        scan_name=scan_name,\n",
    "        attack_strategies=[\n",
    "            AttackStrategy.EASY,\n",
    "            AttackStrategy.MODERATE,\n",
    "            AttackStrategy.DIFFICULT,\n",
    "        ],\n",
    "        output_path=str(output_path),\n",
    "    )\n",
    "    print(f\"Completed {scan_name} -> {output_path}\")\n",
    "except Exception as e:  # noqa: BLE001\n",
    "    print(f\"ERROR during {scan_name}: {e}\")\n",
    "    raise"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
